import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain_openai import OpenAIEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.vectorstores import FAISS
from langchain.memory import VectorStoreRetrieverMemory

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base

memory = ConversationBufferMemory()
memory.save_context(
    {"input":"帮我找一下tomie"},
    {"output":"对不起请问什么是tomie？"}
)
memory.save_context(
    {"input":"tomie是一个培训讲师"},
    {"output":"好的，我知道了。"}
)
memory.save_context(
    {"input":"今天他要讲一门关于RAG的课程"},
    {"output":"好的，我知道了。需要RAG的资料吗？"}
)
memory.save_context(
    {"input":"不需要资料了，谢谢"},
    {"output":"好的，那我就不打扰你了。"}
)

vectorstore = FAISS.from_texts(
    memory.buffer.split("\n"),
    OpenAIEmbeddings()
)
FAISS.save_local(vectorstore,"test_faiss")

# 查找相似内容
# print(FAISS.load_local("test_faiss", OpenAIEmbeddings(model="text-embedding-3-small"),
#                        allow_dangerous_deserialization=True).similarity_search("tomie是什么职业?"))

# 读取存储的记忆并检索
r1 = FAISS.load_local("test_faiss",OpenAIEmbeddings(),allow_dangerous_deserialization=True)
r2 = r1.as_retriever(
    search_kwargs={"k":1}
)
memory2 = VectorStoreRetrieverMemory(
    retriever=r2
)
print(memory2.load_memory_variables({"prompt": "tomie是什么职业?"}))